#!/usr/bin/env python3
"""
增强版跳绳视频分析器 - 结合优化检测算法和视频输出
基于成功的检测算法，生成带轨迹线和计数的分析视频
"""

import cv2
import numpy as np
import os
import json
import argparse
from datetime import datetime, timedelta
from moviepy.editor import VideoFileClip
import pandas as pd

import mediapipe as mp
from mediapipe.python.solutions import pose as mp_pose
from mediapipe.python.solutions import drawing_utils as mp_drawing
from mediapipe.python.solutions import drawing_styles as mp_drawing_styles

class EnhancedVideoAnalyzer:
    def __init__(self, video_path, csv_path, detection_result_path, output_dir=None):
        self.video_path = video_path
        self.csv_path = csv_path
        self.detection_result_path = detection_result_path
        
        # 加载检测结果
        with open(detection_result_path, 'r', encoding='utf-8') as f:
            self.detection_result = json.load(f)
        
        # 加载CSV数据
        try:
            self.csv_data = pd.read_csv(csv_path)
        except Exception as e:
            print(f"CSV读取失败，尝试备用方法: {e}")
            try:
                self.csv_data = pd.read_csv(csv_path, dtype=str)
                # 转换数值列
                numeric_columns = [col for col in self.csv_data.columns if '_x' in col or '_y' in col or '_z' in col or 'visibility' in col or col in ['frame_number', 'timestamp']]
                for col in numeric_columns:
                    if col in self.csv_data.columns:
                        self.csv_data[col] = pd.to_numeric(self.csv_data[col], errors='coerce')
                # 转换布尔列
                bool_columns = ['has_pose']
                for col in bool_columns:
                    if col in self.csv_data.columns:
                        self.csv_data[col] = self.csv_data[col].astype(bool)
            except Exception as e2:
                raise ValueError(f"无法读取CSV文件: {e2}")
        
        self.cap = cv2.VideoCapture(video_path)
        if not self.cap.isOpened():
            raise ValueError(f"无法打开视频文件: {video_path}")
            
        # 获取视频信息
        self.fps = int(self.cap.get(cv2.CAP_PROP_FPS))
        self.total_frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
        self.width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        self.height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        
        # 设置输出目录
        if output_dir is None:
            output_dir = os.path.dirname(video_path)
        self.output_dir = output_dir
        self.base_name = os.path.splitext(os.path.basename(video_path))[0]
        
        # MediaPipe初始化
        self.pose = mp_pose.Pose(
            static_image_mode=False,
            model_complexity=1,
            min_detection_confidence=0.5,
            min_tracking_confidence=0.5
        )
        self.mp_drawing = mp_drawing
        self.mp_pose = mp_pose
        self.mp_drawing_styles = mp_drawing_styles
        
        # 轨迹线历史
        self.trajectory_history = {}
        for i in range(33):
            self.trajectory_history[i] = []
        
        # 跳绳计数显示
        self.current_jump_count = 0
        self.jump_detection_points = self._extract_jump_points()
        
    def _extract_jump_points(self):
        """从检测结果中提取真实的跳跃帧号"""
        jump_frames_data = self.detection_result.get('jump_frames', [])
        jump_frames = []

        if not jump_frames_data:
            print("警告: 'jump_frames' 在检测结果中未找到或为空。")
        elif isinstance(jump_frames_data[0], dict) and 'frame' in jump_frames_data[0]:
            # 新格式: [{"jump_number": 1, "frame": 173}, ...]
            jump_frames = [item['frame'] for item in jump_frames_data]
        elif isinstance(jump_frames_data[0], int):
            # 旧格式: [173, 182, ...]
            jump_frames = jump_frames_data
        else:
            print(f"警告: 未知的 'jump_frames' 格式: {type(jump_frames_data[0])}")

        # 如果jump_frames为空，执行回退逻辑
        if not jump_frames:
            print("执行回退逻辑来估算跳跃帧。")
            total_jumps = self.detection_result.get('detected_jumps', 0)
            total_frames = self.detection_result.get('total_frames', self.total_frames)
            if total_jumps > 0 and total_frames > 0:
                interval = total_frames / total_jumps
                for i in range(total_jumps):
                    frame_num = int(i * interval + interval / 2)
                    if frame_num < total_frames:
                        jump_frames.append(frame_num)
        
        return sorted(jump_frames)
    
    def get_landmarks_from_csv(self, frame_num):
        """从CSV数据中获取指定帧的关键点数据"""
        if frame_num >= len(self.csv_data):
            return None
            
        row = self.csv_data.iloc[frame_num]
        if not row['has_pose']:
            return None
        
        # 构建landmarks数据结构
        landmarks = []
        landmark_names = [
            'nose', 'left_eye_inner', 'left_eye', 'left_eye_outer',
            'right_eye_inner', 'right_eye', 'right_eye_outer',
            'left_ear', 'right_ear', 'mouth_left', 'mouth_right',
            'left_shoulder', 'right_shoulder', 'left_elbow', 'right_elbow',
            'left_wrist', 'right_wrist', 'left_pinky', 'right_pinky',
            'left_index', 'right_index', 'left_thumb', 'right_thumb',
            'left_hip', 'right_hip', 'left_knee', 'right_knee',
            'left_ankle', 'right_ankle', 'left_heel', 'right_heel',
            'left_foot_index', 'right_foot_index'
        ]
        
        for name in landmark_names:
            try:
                x = float(row.get(f'{name}_x', 0))
                y = float(row.get(f'{name}_y', 0))
                z = float(row.get(f'{name}_z', 0))
                visibility = float(row.get(f'{name}_visibility', 0))
            except (ValueError, TypeError):
                x, y, z, visibility = 0.0, 0.0, 0.0, 0.0
            
            # 创建landmark对象
            landmark = type('Landmark', (), {
                'x': x, 'y': y, 'z': z, 'visibility': visibility
            })()
            landmarks.append(landmark)
        
        return landmarks
    
    def draw_trajectory_lines(self, frame, landmarks):
        """绘制轨迹线 - 参考universal_jump_rope_analyzer的实现"""
        if landmarks is None:
            return frame
        
        # 更新轨迹历史
        for i, landmark in enumerate(landmarks):
            try:
                x = int(landmark.x * self.width)
                y = int(landmark.y * self.height)
                # 确保坐标在有效范围内
                if 0 <= x < self.width and 0 <= y < self.height:
                    point = (x, y)
                    self.trajectory_history[i].append(point)
                    if len(self.trajectory_history[i]) > 30:
                        self.trajectory_history[i].pop(0)
            except (AttributeError, ValueError, TypeError):
                continue
        
        # 为每个关键点绘制轨迹线
        for i, history in self.trajectory_history.items():
            if len(history) > 1:
                # 渐变颜色：从蓝色到绿色
                color_ratio = min(1.0, len(history) / 30)
                color = (int(255 * (1 - color_ratio)), int(255 * color_ratio), 0)
                
                # 绘制轨迹线段
                for j in range(len(history) - 1):
                    thickness = max(1, int(color_ratio * 3))
                    cv2.line(frame, history[j], history[j+1], color, thickness)
        
        # 特别为脚部添加更明显的轨迹线
        if len(self.trajectory_history.get(27, [])) > 1 and len(self.trajectory_history.get(28, [])) > 1:
            # 左踝轨迹
            for j in range(len(self.trajectory_history[27]) - 1):
                cv2.line(frame, self.trajectory_history[27][j], self.trajectory_history[27][j+1], (0, 0, 255), 3)
            # 右踝轨迹
            for j in range(len(self.trajectory_history[28]) - 1):
                cv2.line(frame, self.trajectory_history[28][j], self.trajectory_history[28][j+1], (0, 0, 255), 3)
        
        return frame
    
    def update_jump_count(self, frame_num):
        """根据真实的跳跃帧更新跳绳计数"""
        # 计算当前帧之前（包括当前帧）发生了多少次跳跃
        count = sum(1 for jump_frame in self.jump_detection_points if jump_frame <= frame_num)
        self.current_jump_count = count
    
    def draw_analysis_info(self, frame, frame_num):
        """绘制分析信息 - 参考universal_jump_rope_analyzer的界面"""
        h, w = frame.shape[:2]
        
        # 更新跳绳计数
        self.update_jump_count(frame_num)
        
        # 添加黄色标题
        title = "Enhanced Jump Rope Analysis"
        title_font = cv2.FONT_HERSHEY_SIMPLEX
        title_scale = 1.5
        title_thickness = 3
        title_color = (0, 255, 255)  # 黄色
        
        title_size = cv2.getTextSize(title, title_font, title_scale, title_thickness)[0]
        title_x = int((w - title_size[0]) / 2)
        title_y = 60
        
        # 绘制标题背景
        overlay = frame.copy()
        cv2.rectangle(overlay, (0, 0), (w, title_y + 20), (0, 0, 0), -1)
        cv2.addWeighted(overlay, 0.7, frame, 0.3, 0, frame)
        cv2.putText(frame, title, (title_x, title_y), title_font, title_scale, title_color, title_thickness)
        
        # 计算并显示时间
        elapsed_seconds = frame_num / self.fps
        time_str = str(timedelta(seconds=int(elapsed_seconds)))[2:7]  # 格式为 MM:SS
        cv2.putText(frame, time_str, (50, 150), 
                   cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 0), 5)
        cv2.putText(frame, time_str, (50, 150), 
                   cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 255, 255), 3)
        
        # 添加信息面板
        panel_x, panel_y = 10, h - 130
        panel_w, panel_h = 300, 120
        
        # 半透明背景
        overlay = frame.copy()
        cv2.rectangle(overlay, (panel_x, panel_y), (panel_x + panel_w, panel_y + panel_h), (0, 0, 0), -1)
        cv2.addWeighted(overlay, 0.7, frame, 0.3, 0, frame)
        
        # 大号计数显示
        count_font = cv2.FONT_HERSHEY_DUPLEX
        count_scale = 1.2
        count_thickness = 3
        count_str = f"{self.current_jump_count}"
        count_size = cv2.getTextSize(count_str, count_font, count_scale, count_thickness)[0]
        count_x = w - count_size[0] - 50
        count_y = panel_y + count_size[1] + 10
        
        cv2.putText(frame, count_str, (count_x, count_y), 
                   count_font, count_scale, (0, 255, 0), count_thickness)
        
        # 显示跳跃计数和检测方法
        font = cv2.FONT_HERSHEY_DUPLEX
        y_offset = 30
        cv2.putText(frame, f"Jump Count: {self.current_jump_count}", (panel_x + 10, panel_y + y_offset), font, 0.7, (255, 255, 255), 2)
        y_offset += 25
        cv2.putText(frame, f"Method: {self.detection_result['detection_method']}", (panel_x + 10, panel_y + y_offset), font, 0.6, (0, 255, 255), 2)
        y_offset += 25
        cv2.putText(frame, f"Total: {self.detection_result['detected_jumps']}", (panel_x + 10, panel_y + y_offset), font, 0.7, (0, 255, 0), 2)
        
        # 显示进度条
        progress = frame_num / self.total_frames if self.total_frames > 0 else 0
        progress_bar_w = w - 100
        progress_bar_h = 20
        progress_bar_x = 50
        progress_bar_y = h - 30
        
        # 绘制进度条背景
        cv2.rectangle(frame, (progress_bar_x, progress_bar_y), 
                     (progress_bar_x + progress_bar_w, progress_bar_y + progress_bar_h), 
                     (50, 50, 50), -1)
        
        # 绘制进度
        progress_width = int(progress_bar_w * progress)
        cv2.rectangle(frame, (progress_bar_x, progress_bar_y), 
                     (progress_bar_x + progress_width, progress_bar_y + progress_bar_h), 
                     (0, 255, 0), -1)
        
        # 显示进度文本
        progress_text = f"Frame: {frame_num}/{self.total_frames}"
        cv2.putText(frame, progress_text, (progress_bar_x, progress_bar_y - 10), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1)
        
        return frame
    
    def process_video(self, output_path=None):
        """处理整个视频 - 生成带轨迹线和计数的分析视频"""
        print(f"Processing video: {self.video_path}")
        print(f"Video info: {self.width}x{self.height}, {self.fps}fps, {self.total_frames} frames")
        print(f"Detection result: {self.detection_result['detected_jumps']} jumps using {self.detection_result['detection_method']}")
        
        # 准备输出路径
        if output_path is None:
            output_path = os.path.join(self.output_dir, f"{self.base_name}_enhanced_analyzed.mp4")
        
        # 创建临时文件路径（用于后续音频合并）
        temp_video_path = output_path.replace('.mp4', '_temp.mp4')
        
        # 创建视频写入器 - 使用与参考文件相同的编码器
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(temp_video_path, fourcc, float(self.fps), (self.width, self.height))
            
        if not out.isOpened():
            raise ValueError("无法创建视频写入器")
        
        frame_count = 0
        
        while True:
            ret, frame = self.cap.read()
            if not ret:
                break
            
            # 获取当前帧的关键点数据
            landmarks = self.get_landmarks_from_csv(frame_count)
            
            # 绘制关键点（简化版本）
            if landmarks:
                # 直接绘制关键点，避免MediaPipe兼容性问题
                for i, landmark in enumerate(landmarks):
                    try:
                        if landmark.visibility > 0.5:
                            x = int(landmark.x * self.width)
                            y = int(landmark.y * self.height)
                            # 确保坐标在有效范围内
                            if 0 <= x < self.width and 0 <= y < self.height:
                                cv2.circle(frame, (x, y), 3, (0, 255, 0), -1)
                    except (AttributeError, ValueError, TypeError):
                        continue
            
            # 绘制轨迹线
            frame = self.draw_trajectory_lines(frame, landmarks)
            
            # 绘制分析信息
            frame = self.draw_analysis_info(frame, frame_count)
            
            # 写入视频
            out.write(frame)
            
            frame_count += 1
            
            # 显示进度
            if frame_count % 100 == 0:
                print(f"Progress: {frame_count}/{self.total_frames} frames, current count: {self.current_jump_count}")
        
        # 释放资源
        self.cap.release()
        out.release()
        
        print(f"Video processing complete, adding audio...")
        
        # 添加音频
        try:
            # 读取原始视频和临时视频
            original_video = VideoFileClip(self.video_path)
            processed_video = VideoFileClip(temp_video_path)
            
            # 将原始音频添加到处理后的视频
            final_video = processed_video.set_audio(original_video.audio)
            
            # 写入最终视频文件
            final_video.write_videofile(
                output_path,
                codec='libx264',
                audio_codec='aac',
                temp_audiofile='temp-audio.m4a',
                remove_temp=True
            )
            
            # 清理资源
            original_video.close()
            processed_video.close()
            final_video.close()
            
            # 删除临时文件
            import time
            time.sleep(1)
            if os.path.exists(temp_video_path):
                os.remove(temp_video_path)
            
            print(f"✅ Audio added successfully!")
            print(f"Final output file: {output_path}")
            
        except Exception as e:
            print(f"⚠️ Audio processing failed: {e}")
            print("Keeping audio-free version")
            
            # 如果音频处理失败，将临时文件重命名为最终文件
            if os.path.exists(temp_video_path):
                if os.path.exists(output_path):
                    os.remove(output_path)
                os.rename(temp_video_path, output_path)
                print(f"Saved audio-free version: {output_path}")
        
        print(f"Processing complete!")
        print(f"Final jump count: {self.current_jump_count}")
        print(f"Detection accuracy: {self.detection_result['detected_jumps']} jumps detected")

def main():
    parser = argparse.ArgumentParser(description='Enhanced Jump Rope Video Analyzer')
    parser.add_argument('video_path', help='Path to the input video file')
    parser.add_argument('csv_path', help='Path to the CSV frame data file')
    parser.add_argument('detection_result_path', help='Path to the detection result JSON file')
    parser.add_argument('--output', help='Output video path (optional)')
    
    args = parser.parse_args()
    
    # 创建分析器并处理视频
    analyzer = EnhancedVideoAnalyzer(
        args.video_path, 
        args.csv_path, 
        args.detection_result_path
    )
    analyzer.process_video(args.output)

if __name__ == "__main__":
    main()